import os
from flask_restful import Resource
from flask import request
from common.resp import ajax_success, ajax_error, ajax_data, ajax_page
from core import db, app
from datetime import datetime

class ModelsModel(db.Model):
    __tablename__ = "sys_models"

    id = db.Column(db.Integer, autoincrement=True, primary_key=True)
    code = db.Column(db.String(50), nullable=False, index=True, unique=True)
    name = db.Column(db.String(50), nullable=False, index=True)
    remark = db.Column(db.Text, nullable=True)
    path = db.Column(db.String(255), nullable=False)
    version = db.Column(db.String(10), nullable=False)  # yolov8n, yolov8s, yolov8m, yolov8l, yolov8x
    type = db.Column(db.String(50), nullable=False, default="detect")  # detect, segment, classify, pose
    status = db.Column(db.String(50), nullable=False, default="1")  # 1: 启用, 0: 停用
    create_time = db.Column(db.DateTime, default=datetime.now)
    updated_time = db.Column(db.DateTime, default=datetime.now, onupdate=datetime.now)
    
    def to_json(self):
        """将模型对象转换为字典，支持JSON序列化"""
        return {
            'id': self.id,
            'code': self.code,
            'name': self.name,
            'remark': self.remark,
            'path': self.path,
            'version': self.version,
            'type': self.type,
            'status': self.status,
            'create_time': self.create_time.strftime('%Y-%m-%d %H:%M:%S') if self.create_time else None,
            'updated_time': self.updated_time.strftime('%Y-%m-%d %H:%M:%S') if self.updated_time else None
        }


class ModelsResource(Resource):
    def get(self):
        query = request.args;
        filters = []
        if query.get('name'):
            filters.append(ModelsModel.name.like('%' + query['name'] + '%'))
        if query.get('code'):
            filters.append(ModelsModel.code.like('%' + query['code'] + '%'))
        if query.get('type'):
            filters.append(ModelsModel.type == query['type'])
        if query.get('status'):
            filters.append(ModelsModel.status == query['status'])
        if query.get('version'):
            filters.append(ModelsModel.version == query['version'])

        paginate = ModelsModel.query.filter(
            *filters).paginate(
            page=int(query.get('page', 1)), per_page=int(query.get('limit', 10)), error_out=False)

        # 将模型对象转换为字典列表以支持JSON序列化
        items_dict = [item.to_json() for item in paginate.items]
        return ajax_page(items_dict, paginate.total)

    def post(self):
        data = request.get_json()
        model = ModelsModel(
            code = data.get('code'),
            name = data.get('name'),
            remark = data.get('remark'),
            type = data.get('type', 'detect'),
            version = data.get('version', 'yolov8n'),
            status = data.get('status', '1'),
            path = data.get('path', data.get('code') + '/default')
        )

        modelPath = os.path.join(app.config['MODEL_PATH'], model.path)
        if not os.path.exists(modelPath):
            os.makedirs(modelPath)

        db.session.add(model)
        db.session.commit()
        return ajax_success('模型创建成功')

    def put(self):
        data = request.get_json()
        if not data.get('id'):
            return ajax_error('参数错误')

        model = ModelsModel.query.get_or_404(data.get('id'))
        if not model:
            return ajax_error('模型不存在')

        if data.get('name'):
            model.name = data.get('name', model.name)
        if data.get('status'):
            model.status = data.get('status', model.status)
        if data.get('remark'):
            model.remark = data.get('remark', model.remark)    
        if data.get('type'):
            model.type = data.get('type', model.type)
        if data.get('version'):
            model.version = data.get('version', model.version)

        db.session.add(model)
        db.session.commit()
        db.session.refresh(model)
        return ajax_success('模型更新成功')



class ModelsDetail(Resource):
    def get(self, id):
        if id:
            model = ModelsModel.query.get_or_404(id)
            return ajax_data(model.to_json())
        else:
            return ajax_error('参数错误')

    def delete(self, id):
        from modules.training import TrainingModel
        t = TrainingModel.query.filter_by(model_id=id).first()
        if t:
            return ajax_error('删除模型前请先删除训练任务')

        # 先删除目录
        model = ModelsModel.query.get_or_404(id)
        modelPath = os.path.join(app.config['MODEL_PATH'], model.path)
        if os.path.exists(modelPath):
            os.rmdir(modelPath)

        #再删除数据库
        db.session.delete(model)
        db.session.commit()
        return ajax_success('模型删除成功')